import time
from threading import Thread
import cv2
import ffmpegcv
import torch
from flask import Flask
from ultralytics import YOLO
import logging

population = 0
stop_monitor = False
app = Flask("ScenicSpotPopulationDetector")
log = logging.getLogger('werkzeug')
log.setLevel(logging.ERROR)


def start_server_thread(port=8888):
    thread_flask = Thread(target=start_flask, args=(port,))
    thread_flask.daemon = True
    thread_flask.start()


def detect(video_path, detection_listener, interval, frame_interval):
    global population
    global stop_monitor

    model = YOLO("./yolov8_face.pt")

    # overrides = {'conf': 0.1, 'iou': 0.6, 'agnostic_nms': False, 'max_det': 1000, 'device': 'cpu', 'task': 'detect'}

    # model.overrides = overrides

    capture = ffmpegcv.VideoCapture(video_path)

    stop_monitor = False

    now_frame = 0

    speedtest_epoch = 0
    speedtest_start = time.time()

    while not stop_monitor and speedtest_epoch < 10:
        ret, frame = capture.read()

        if frame.shape != (640, 640, 3):
            raise Exception("Invalid vide size, only videos with 640x640 are supported.")

        if not ret:
            break

        now_frame += 1
        if now_frame % frame_interval != 0:
            continue

        results = model(frame.astype(float), verbose=False)

        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        for pos in results[0].boxes.xyxy:
            x1, y1, x2, y2 = map(int, pos)
            cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
        result_number = len(results[0])
        population = result_number

        detection_listener(frame, result_number)

        time.sleep(interval)

        speedtest_epoch += 1

    capture.release()

    speedtest_end = time.time()
    print("Time for 10 epochs: ", speedtest_end - speedtest_start)


def start_detect_thread(video_path, detection_listener, interval, frame_interval):
    # Load the backend.
    torch.ops.load_library("build/librvv_pytorch.so")
    Thread(target=detect, args=(video_path, detection_listener, interval, frame_interval)).start()


@app.route('/population', methods=['GET'])
def get_data():
    return str(population)


def start_flask(port):
    app.run(port=port)


def stop_monitoring():
    global stop_monitor
    stop_monitor = True
